import cv2
import numpy as np
from scipy.interpolate import interp1d
import random
import os


def add_random_lines_in_segments(image_path, yolo_txt_path):
    # 读取图片
    image = cv2.imread(image_path)
    if image is None:
        print(f"无法读取图片: {image_path}")
        return

    height, width = image.shape[:2]

    try:
        with open(yolo_txt_path, 'r') as f:
            lines = f.readlines()
    except FileNotFoundError:
        print(f"未找到标注文件: {yolo_txt_path}")
        return

    # 创建一个与原图像相同大小的空白覆盖层
    overlay = image.copy()

    for line in lines:
        parts = line.strip().split()
        if len(parts) == 0:
            continue
        class_id = int(parts[0])
        polygon_points = [float(x) for x in parts[1:]]

        # 将归一化坐标转换为实际坐标
        points = []
        for i in range(0, len(polygon_points), 2):
            x = int(polygon_points[i] * width)
            y = int(polygon_points[i + 1] * height)
            points.append([x, y])
        points = np.array(points, dtype=np.int32)

        # 随机确定线条的数量
        num_lines = random.randint(2, 3)
        for _ in range(num_lines):
            # 在多边形内部随机选择起点和终点
            while True:
                start_x = random.randint(points[:, 0].min(), points[:, 0].max())
                start_y = random.randint(points[:, 1].min(), points[:, 1].max())
                if cv2.pointPolygonTest(points, (start_x, start_y), False) >= 0:
                    break
            while True:
                end_x = random.randint(points[:, 0].min(), points[:, 0].max())
                end_y = random.randint(points[:, 1].min(), points[:, 1].max())
                if cv2.pointPolygonTest(points, (end_x, end_y), False) >= 0:
                    break

            # 随机选择线条颜色（这里使用 BGR 格式）
            line_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
            # 随机选择线条粗细
            line_thickness = random.randint(4, 8)

            # 在覆盖层上绘制线条
            cv2.line(overlay, (start_x, start_y), (end_x, end_y), line_color, line_thickness)

    # 随机选择不透明度（范围 0.5 - 1）
    alpha = random.uniform(0.5, 1)
    # 将覆盖层与原图像按不透明度混合
    cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image)

    return image


def add_random_curve(image):
    height, width = image.shape[:2]
    # 随机确定曲线的数量，范围为 2 到 3 条
    num_curves = random.randint(2, 3)

    for _ in range(num_curves):
        # 随机确定控制点的数量
        num_control_points = random.randint(4, 8)
        # 随机生成控制点的 x 坐标
        control_points_x = sorted([random.randint(0, width - 1) for _ in range(num_control_points)])
        # 随机生成控制点的 y 坐标
        control_points_y = [random.randint(0, height - 1) for _ in range(num_control_points)]

        # 创建插值函数
        f = interp1d(control_points_x, control_points_y, kind='linear')

        # 生成曲线上的更多点
        curve_x = np.linspace(min(control_points_x), max(control_points_x), num=200)
        curve_y = f(curve_x)

        # 将曲线上的点转换为整数坐标
        curve_points = np.column_stack((curve_x.astype(int), curve_y.astype(int)))

        # 随机选择曲线颜色（这里使用 BGR 格式）
        curve_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
        # 随机选择曲线粗细
        curve_thickness = random.randint(4, 8)
        # 随机选择曲线不透明度（范围 0.5 - 1）
        alpha = random.uniform(0.5, 1)

        # 创建一个与原图像相同大小的空白图像，用于绘制曲线
        overlay = image.copy()
        cv2.polylines(overlay, [curve_points], isClosed=False, color=curve_color, thickness=curve_thickness)

        # 将绘制了曲线的图像与原图像按透明度混合
        cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image)

    return image


def cover(source_dir, yolo_txt_dir, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    # 遍历源目录下的所有图像文件
    for filename in os.listdir(source_dir):
        if filename.endswith(('.png', '.jpg', '.jpeg')):
            # 构建图像文件的完整路径
            image_path = os.path.join(source_dir, filename)
            base_name = os.path.splitext(filename)[0]
            yolo_txt_path = os.path.join(yolo_txt_dir, f"{base_name}.txt")
            output_image = add_random_lines_in_segments(image_path, yolo_txt_path)
            # 添加随机曲线
            image_with_curve = add_random_curve(output_image)
            # 构建输出图像的完整路径
            output_path = os.path.join(output_dir, filename)
            # 保存添加曲线后的图像
            cv2.imwrite(output_path, image_with_curve)
            print(f"已处理 {filename} 并保存到 {output_path}")


if __name__ == '__main__':
    source_dir = "../test/val/images"
    yolo_txt_dir = "../test/val/labels"
    output_dir = "../test/val/images"
    cover(source_dir, yolo_txt_dir, output_dir)

    source_dir = "../test/train/images"
    yolo_txt_dir = "../test/train/labels"
    output_dir = "../test/train/images"
    cover(source_dir, yolo_txt_dir, output_dir)
